#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

from pymongo import MongoClient
import re


# 连接MongoDB
def InitMongoDB(MONGODB_HOST, MONGODB_PORT, MONGODB_DB_NAME, MONGODB_COLLECTION_NAME):
    client = MongoClient(MONGODB_HOST, MONGODB_PORT)
    db = client[MONGODB_DB_NAME]
    collection = db[MONGODB_COLLECTION_NAME]
    return client, collection


# 关闭MongoDB
def CloseMongoDB(MongoDB_Connection):
    MongoDB_Connection.close()


# 第一步导入imsg中所有Species信息
def ImportSpeciesFromIMSG():
    client_IMSG, collection_IMSG = InitMongoDB("10.188.188.22", 27017, "Pre_Data", "imsg")
    client_eLMSG, collection_eLMSG = InitMongoDB("10.188.188.22", 27017, "eLMSG", "eLMSG")

    results = collection_IMSG.find({"Rank": "Species"})
    for result in results:
        result.pop("_id")
        collection_eLMSG.insert_one(result)
    CloseMongoDB(client_IMSG)
    CloseMongoDB(client_eLMSG)


# 第二步从Taxonomy导入其他Rank,并查询imsg中是否有对应信息,一并导入
def ImportOtherRanksFromTaxonomy():
    MicrobiomeRoots = set(["2157", "2", "4751", "10239", "12884", "151659", "408169"])

    client_IMSG, collection_IMSG = InitMongoDB("10.188.188.22", 27017, "Pre_Data", "imsg")
    client_eLMSG, collection_eLMSG = InitMongoDB("10.188.188.22", 27017, "eLMSG", "eLMSG")
    client_Taxonomy, collection_Taxonomy = InitMongoDB("10.188.188.22", 27017, "Dictionary", "taxonomy_20170421")

    for result in collection_eLMSG.find({"Rank": "Species"}):
        if "TaxonID" in result:
            RecursionTaxon(result["TaxonID"], collection_Taxonomy, collection_IMSG, collection_eLMSG, MicrobiomeRoots)

    CloseMongoDB(client_IMSG)
    CloseMongoDB(client_eLMSG)
    CloseMongoDB(client_Taxonomy)


def RecursionTaxon(TaxonID, collection_Taxonomy, collection_IMSG, collection_eLMSG, MicrobiomeRoots):
    parentTaxonID = collection_Taxonomy.find_one({"dms_id": TaxonID})["parent_tax_id"]["id"]

    if collection_eLMSG.find_one({"TaxonID": parentTaxonID}) == None:
        InIMSG = collection_IMSG.find_one({"TaxonID": parentTaxonID})
        if InIMSG != None:
            InIMSG.pop("_id")
            collection_eLMSG.insert_one(InIMSG)
        else:
            InTaxonomy = collection_Taxonomy.find_one({"tax_id": parentTaxonID})
            result = {}
            for key in InTaxonomy.keys():
                if key.startswith("names_"):
                    result[key] = InTaxonomy[key]
                if key.startswith("division_"):
                    result[key] = InTaxonomy[key]
                if key.startswith("genetic_"):
                    result[key] = InTaxonomy[key]
                if key == "rank":
                    result["Rank"] = InTaxonomy["rank"]
                if key == "tax_id":
                    result["TaxonID"] = InTaxonomy["tax_id"]
                if key == "parent_tax_id":
                    result["parent_tax_id"] = InTaxonomy["parent_tax_id"]["id"]
            collection_eLMSG.insert_one(result)
        if parentTaxonID not in MicrobiomeRoots:
            RecursionTaxon(parentTaxonID, collection_Taxonomy, collection_IMSG, collection_eLMSG, MicrobiomeRoots)


# 第三步导入Bacdive中的信息
def ImportBacdive():
    client_Bacdive, collection_Bacdive = InitMongoDB("10.188.188.22", 27017, "Pre_Data", "bacdive")
    client_eLMSG, collection_eLMSG = InitMongoDB("10.188.188.22", 27017, "eLMSG", "eLMSG")

    for result in collection_eLMSG.find({"Rank": "Species"}):
        if "TaxonID" in result:
            scientific_name = result["names_scientific_name"]
            count = collection_Bacdive.count({"taxonomy_name.strains_tax_PNU.species": scientific_name,
                                              "taxonomy_name.strains_tax_PNU.is_type_strain": True})
            if count == 1:
                bacdive_record = collection_Bacdive.find_one({"taxonomy_name.strains_tax_PNU.species": scientific_name,
                                                              "taxonomy_name.strains_tax_PNU.is_type_strain": True})
                collection_eLMSG.update_one({"names_scientific_name": scientific_name},
                                            {"$set": {"bacdive_backup": bacdive_record}})
    CloseMongoDB(client_eLMSG)
    CloseMongoDB(client_Bacdive)


def GenerateCitationReference():
    client_eLMSG, collection_eLMSG = InitMongoDB("10.188.188.22", 27017, "eLMSG", "eLMSG")

    for result in collection_eLMSG.find({"citaion_reference": {"$exists": 1}}, {"citaion_reference": 1}):
        reference = []
        for line in result["citaion_reference"]:
            reference.append(GetReference(line))
        collection_eLMSG.update_one({"_id": result["_id"]}, {"$set": {"CitationReference": reference}})
    CloseMongoDB(client_eLMSG)


def GenerateMergedReference():
    client_eLMSG, collection_eLMSG = InitMongoDB("10.188.188.22", 27017, "eLMSG", "eLMSG")

    for result in collection_eLMSG.find({"Reference": {"$exists": 1}, "CitationReference": {"$exists": 1}},
                                        {"Reference": 1, "CitationReference": 1}):
        merged_reference = MergeReference(result["Reference"], result["CitationReference"])
        collection_eLMSG.update_one({"_id": result["_id"]}, {"$set": {"MergedReference": merged_reference}})
    for result in collection_eLMSG.find({"Reference": {"$exists": 1}, "CitationReference": {"$exists": 0}},
                                        {"Reference": 1}):
        merged_reference = result["Reference"]
        collection_eLMSG.update_one({"_id": result["_id"]}, {"$set": {"MergedReference": merged_reference}})
    for result in collection_eLMSG.find({"Reference": {"$exists": 0}, "CitationReference": {"$exists": 1}},
                                        {"CitationReference": 1}):
        merged_reference = result["CitationReference"]
        collection_eLMSG.update_one({"_id": result["_id"]}, {"$set": {"MergedReference": merged_reference}})
    CloseMongoDB(client_eLMSG)


def GetReference(line):
    reference = {}
    regex = "\(((PM|DO).*)\)"
    prog = re.compile(regex)
    result = prog.search(line)
    if result:
        ids = result.group(1)
        for id in re.split(r'[\s ]', ids):
            if id.startswith("PMID"):
                reference["PMID"] = id[5:]
            elif id.startswith("DOI"):
                reference["DOI"] = id[4:]
            elif id.startswith("PMCID"):
                reference["PMCID"] = id[6:]
        txt = line.replace("(" + ids + ")", "", -1).strip()
        if txt != "":
            reference["txt"] = txt
    else:
        txt = line.strip()
        if txt != "":
            reference["txt"] = txt
    return reference


def MergeReference(refset1, refset2):
    refset3 = refset1 + refset2
    for i in range(0, len(refset3)):
        if i >= len(refset1):
            break
        else:
            for j in range(i + 1, len(refset3)):
                if "DOI" in refset3[i]:
                    if "DOI" in refset3[j]:
                        if refset3[i]["DOI"] == refset3[j]["DOI"]:
                            if "PMID" in refset3[j]:
                                refset3[i]["PMID"] = refset3[j]["PMID"]
                            if "PMCID" in refset3[j]:
                                refset3[i]["PMCID"] = refset3[j]["PMCID"]
                            if "txt" in refset3[j]:
                                refset3[i]["txt"] = refset3[j]["txt"]
                            refset3.pop(j)
                            break
                if "PMID" in refset3[i]:
                    if "PMID" in refset3[j]:
                        if refset3[i]["PMID"] == refset3[j]["PMID"]:
                            if "DOI" in refset3[j]:
                                refset3[i]["DOI"] = refset3[j]["DOI"]
                            if "PMCID" in refset3[j]:
                                refset3[i]["PMCID"] = refset3[j]["PMCID"]
                            if "txt" in refset3[j]:
                                refset3[i]["txt"] = refset3[j]["txt"]
                            refset3.pop(j)
                            break
                if "PMCID" in refset3[i]:
                    if "PMCID" in refset3[j]:
                        if refset3[i]["PMCID"] == refset3[j]["PMCID"]:
                            if "DOI" in refset3[j]:
                                refset3[i]["DOI"] = refset3[j]["DOI"]
                            if "PMID" in refset3[j]:
                                refset3[i]["PMID"] = refset3[j]["PMID"]
                            if "txt" in refset3[j]:
                                refset3[i]["txt"] = refset3[j]["txt"]
                            refset3.pop(j)
                            break
    return refset3


# 生成索引
def GenerateIndex():
    import pymongo
    client_eLMSG, collection_eLMSG = InitMongoDB("10.188.188.22", 27017, "eLMSG", "eLMSG")

    collection_eLMSG.create_index([("$**", pymongo.TEXT)], background=True)

    CloseMongoDB(client_eLMSG)


def FullfillLMSG_ID(startID):
    client_eLMSG, collection_eLMSG = InitMongoDB("10.188.188.22", 27017, "eLMSG", "eLMSG")

    for result in collection_eLMSG.find({"LMSG_ID": {"$exists": 0}}, {"_id": 1}):
        startID += 1
        lmsg_id = "MSG" + str(startID).zfill(6)
        collection_eLMSG.update_one({"_id": result["_id"]}, {"$set": {"LMSG_ID": lmsg_id}})

    CloseMongoDB(client_eLMSG)


def ChangeAuthorityFieldType():
    client_eLMSG, collection_eLMSG = InitMongoDB("10.188.188.22", 27017, "eLMSG", "eLMSG")

    for result in collection_eLMSG.find({"names_authority": {"$exists": 1}}, {"names_authority": 1}):
        if isinstance(result["names_authority"], str):
            name = result["names_authority"]
            id = result["_id"]
            collection_eLMSG.update_one({"_id": id}, {"$set": {"names_authority": [name]}})

    CloseMongoDB(client_eLMSG)


if __name__ == '__main__':
    ImportSpeciesFromIMSG()
    ImportOtherRanksFromTaxonomy()
    FullfillLMSG_ID(6297)
    ChangeAuthorityFieldType()
    GenerateCitationReference()
    GenerateMergedReference()
    ImportBacdive()
    # GenerateIndex()
